package cn.edu.nju.ws.sview.reasoning;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;

/**
 * 
 * @author Gong Saisai
 * working memeory
 *
 */
public class WorkingMemory {
	private HashMap<Predicate,HashSet<GroundAtom>> predInstantiations;
    private HashMap<Predicate,HashMap<Constant,HashSet<GroundAtom>>> firstArgIndex ;
    private HashMap<Predicate,HashMap<Constant,HashSet<GroundAtom>>> secondArgIndex ;
    private HashMap<Predicate,HashMap<Constant,HashMap<Constant,HashSet<GroundAtom>>>> first_secondArgIndex ;
   
    public WorkingMemory(){
    	this.predInstantiations = new HashMap<Predicate,HashSet<GroundAtom>>();
    	this.firstArgIndex = new HashMap<Predicate,HashMap<Constant,HashSet<GroundAtom>>>();
    	this.secondArgIndex = new HashMap<Predicate,HashMap<Constant,HashSet<GroundAtom>>>();
    	this.first_secondArgIndex = new HashMap<Predicate,HashMap<Constant,HashMap<Constant,HashSet<GroundAtom>>>>();
    }
    
    //reset 
    public void Reset(){
    	this.predInstantiations.clear();
    	this.first_secondArgIndex.clear();
    	this.firstArgIndex.clear();
    	this.secondArgIndex.clear();
    }
    
    /**
      * get the facts of pred
      * @param pred
      * @return  empty if none
    */
    public HashSet<GroundAtom> getRelation(Predicate pred){
    	 HashSet<GroundAtom> relation = this.predInstantiations.get(pred) ;
    	 if(relation==null){
    		 relation = new  HashSet<GroundAtom>();
    	 }
    	 return relation;
    }
    
    public  HashMap<Predicate,HashSet<GroundAtom>> getAllRelation(){
    	return this.predInstantiations;
    }
   
    /**
     *add a fact
     * @param fact
     * @return. whether add successfully
     */
    public boolean add(GroundAtom fact){
    	boolean flag =false;
    	Predicate pred = fact.getPredicate();
    	HashSet<GroundAtom> relation = this.predInstantiations.get(pred);
    	if(relation==null){
    		relation = new HashSet<GroundAtom>();
     		this.predInstantiations.put(pred, relation);
    	}
        flag = relation.add(fact) ;
        if(flag)
           updateIndex(fact) ;
        return flag;
    }
    
    /**
     * @param set
     * @return.ֻ whether add successfully
     */
    public boolean addAll(HashSet<GroundAtom> set){
    	boolean flag =false;
    	for(GroundAtom fact:set){
    	    boolean f = add(fact) ;
    	    if(f)
    	    	flag = true;
    	}
    	return flag;
    }

    /**
     * update the index when add an fact
     * @param fact
     */
    public void updateIndex(GroundAtom fact){
    	Predicate pred = fact.getPredicate();
		int arity = pred.getArity();
		if (arity == 1||arity==2||arity>=3) {
			Constant constant = fact.get(0) ;
			HashMap<Constant,HashSet<GroundAtom>> map = this.firstArgIndex.get(pred) ;
			if(map==null){
				map = new HashMap<Constant,HashSet<GroundAtom>>();
				this.firstArgIndex.put(pred, map) ;
			}
			HashSet<GroundAtom> set = map.get(constant) ;
			if(set==null){
				set = new HashSet<GroundAtom>();
				map.put(constant, set) ;
			}
			set.add(fact) ;
		}
		if (arity == 2) {
			Constant firstTerm = fact.get(0) ;
			Constant secondTerm = fact.get(1) ;
			HashMap<Constant,HashSet<GroundAtom>> map = this.secondArgIndex.get(pred) ;
			if(map==null){
				map = new HashMap<Constant,HashSet<GroundAtom>>();
				this.secondArgIndex.put(pred, map) ;
			}
			HashSet<GroundAtom> set = map.get(secondTerm) ;
			if(set==null){
				set = new HashSet<GroundAtom>();
				map.put(secondTerm, set) ;
			}
			set.add(fact) ;
			HashMap<Constant,HashMap<Constant,HashSet<GroundAtom>>> doublemap = this.first_secondArgIndex.get(pred) ;
			if(doublemap==null){
				doublemap = new HashMap<Constant,HashMap<Constant,HashSet<GroundAtom>>>();
				this.first_secondArgIndex.put(pred, doublemap) ;
			}
			HashMap<Constant,HashSet<GroundAtom>> firstArgmap = doublemap.get(firstTerm);
			if(firstArgmap==null){
				firstArgmap = new HashMap<Constant,HashSet<GroundAtom>>();
				doublemap.put(firstTerm, firstArgmap) ;
			}
			HashSet<GroundAtom> secondArgSet = firstArgmap.get(secondTerm) ;
			if(secondArgSet==null){
				secondArgSet = new HashSet<GroundAtom>();
				firstArgmap.put(secondTerm, secondArgSet) ;
			}
			secondArgSet.add(fact) ;
		}
    }
    
   /**
    *get the facts according to the reasoning pattern
    * @param pred
    * @param key. values: null,firstarg,secondarg,botharg
    * @param equalTerms.  equal terms inAtom pattern 
    * @param prevdelta.Seminaive delta
    * @param incrtFlag. 
    * @return
    */
    public HashSet<GroundAtom> getTupleForQuery(Predicate pred,QueryArg key,ArrayList<ArrayList<Integer>> equalTerms,HashMap<Predicate,HashSet<GroundAtom>> prevdelta,boolean incrtFlag){
    	HashSet<GroundAtom> result = new HashSet<GroundAtom>();
    	HashSet<GroundAtom> tempResult = new HashSet<GroundAtom>();
    	// all the arguments not bound(variables)
    	if(key==null){
    		tempResult.addAll(this.getRelation(pred)) ;
    	}else{
    		if(key instanceof FirstArg){
				FirstArg firstArg = (FirstArg) key;
				Constant constant = firstArg.getFirstArg();
				HashMap<Constant, HashSet<GroundAtom>> map = this.firstArgIndex
						.get(pred);
				
				if (map != null) {
					HashSet<GroundAtom> set = map.get(constant);
					if (set != null) {
						tempResult = set;
					}
				}
    		}else if(key instanceof SecondArg){
    			SecondArg secondArg = (SecondArg) key;
				Constant constant = secondArg.getSecondArg();
				HashMap<Constant, HashSet<GroundAtom>> map = this.secondArgIndex
						.get(pred);
				if (map != null) {
					HashSet<GroundAtom> set = map.get(constant);
					if (set != null) {
						tempResult = set;
					}
				}
    		}else{
    			BothArg bothArg = (BothArg) key;
    			Constant firstConst = bothArg.getFirstArg();
				Constant secondConst = bothArg.getSecondArg();
				HashMap<Constant,HashMap<Constant,HashSet<GroundAtom>>> doubleMap = this.first_secondArgIndex
						.get(pred);
				if (doubleMap!=null) {
					HashMap<Constant,HashSet<GroundAtom>> map = doubleMap.get(firstConst);
					if (map  != null) {
						HashSet<GroundAtom> set = map.get(secondConst);
						if(set!=null){
							tempResult = set;
						}
					}
				}
    		}
    	}
    	//for p(?x,?x), arg1 and arg2 must be equal
		if(equalTerms!=null&&equalTerms.size()>0){
			for(GroundAtom fact:tempResult){
				boolean flag = true;
				for(ArrayList<Integer> equalPos:equalTerms){
					HashSet<Constant> set = new HashSet<Constant>();
					for(Integer pos:equalPos){
						set.add(fact.getTerms().get(pos)) ;
					}
					if(set.size()>1){
						flag = false;
						break;
					}
				}
				if(flag)
					result.add(fact) ;
			}
		}else{
			result = tempResult;
		}
    	if(incrtFlag){
    		HashSet<GroundAtom> mResult = new HashSet<GroundAtom>();
    		HashSet<GroundAtom> mRelation = prevdelta.get(pred) ;
			if (result.size() < mRelation.size()) {
				for (GroundAtom fact : result) {
					if (mRelation.contains(fact)) {
						mResult.add(fact);
					}
				}
			}else{
				for (GroundAtom fact : mRelation) {
					if (result.contains(fact)) {
						mResult.add(fact);
					}
				}				
			}
    		return mResult;
    	}
    	return result;
    }
    
    /**
     * check whether exist the fact. 
     * @param fact
     * @return null if none
     */
    public GroundAtom checkExistAtom(GroundAtom fact){
        Predicate pred = fact.getPredicate();
        ArrayList<Integer> constArray = fact.getPattern().getConstIndices();
        ArrayList<Constant> argArray = fact.getPattern().getConstTerms();
        QueryArg key = null;
        if(constArray.size()==2){
        	 key = new BothArg(argArray.get(0),argArray.get(1)) ;
        }else {
        	if(constArray.get(0).equals(new Integer(0))){
        		key = new FirstArg(argArray.get(0));
        	}else{
        		key = new SecondArg(argArray.get(1));
        	}
        }
    	HashSet<GroundAtom> mSet = getTupleForQuery(pred,key,fact.getPattern().getEqualTerms(),null,false);
        if(mSet.size()==0){
     	   return null;
        }
        if(mSet.size()>1){
        	throw new IllegalArgumentException("Only fact can be used here!");
        }
        return fact;
    }
    
    public String toString(){
    	StringBuffer sb = new StringBuffer();
    	Iterator<Predicate> itor = this.predInstantiations.keySet().iterator();
    	while(itor.hasNext()){
    		Predicate pred = itor.next();
    		HashSet<GroundAtom> relation = this.predInstantiations.get(pred) ;
    		sb.append(pred+" : ");
    		sb.append(relation);
    		sb.append("\n");
    	}
    	return sb.toString();
    }
}
